/**
 * Copyright (c) 2020 Raspberry Pi (Trading) Ltd.
 *
 * SPDX-License-Identifier: BSD-3-Clause
 */
#include <string.h>
#include <stdlib.h>
#include "hardware.h"
#include "pca9557.h"
#include "arduino.h"
#include "psram.pio.h"
#include "psram.h"
#include "pico/stdlib.h"
#include "hardware/dma.h"
#include "hardware/clocks.h"
#include "hardware/vreg.h"
#include "log.h"
#include "scope_helper.h"
#include "httpd_cgi.h"

#define PSRAM_PIO_INDEX                 1
#define PSRAM_SM                        0

#define PSRAM_PIO                       INDEX_TO_PIO(PSRAM_PIO_INDEX)
#define PSRAM_DMA_TX_REQ                INDEX_TO_TX_REQ(PSRAM_PIO_INDEX, PSRAM_SM)
#define PSRAM_DMA_RX_REQ                INDEX_TO_RX_REQ(PSRAM_PIO_INDEX, PSRAM_SM)

#define PSRAM_SIZE                      (PSRAM_BITS * 2 * 1024 * 1024)
#define PSRAM_ADDR_MASK                 ((PSRAM_SIZE / 2) - 1)

#define PSRAM_WRITE_CMD                 0x8833
#define PSRAM_READ_CMD                  0xbbee

static int psram_send_pc;
static int psram_send_nostop_pc;
static int psram_recv_pc;
static pio_sm_config g_psram_pio_cfg_tx;
static pio_sm_config g_psram_pio_cfg_rx;
static uint32_t g_psram_dma_send_chan;
static uint32_t g_psram_dma_recv_chan;

static void psram_send(uint32_t *buffer, uint32_t bytes)
{
    pio_sm_set_enabled(PSRAM_PIO, PSRAM_SM, false);
    pio_sm_set_enabled(PSRAM_PIO, PSRAM_SM, false);
    dma_channel_abort(g_psram_dma_send_chan);
    dma_channel_abort(g_psram_dma_recv_chan);
    pio_sm_set_consecutive_pindirs(PSRAM_PIO, PSRAM_SM, PIN_SRAM_D0, PSRAM_BITS, true);
    pio_sm_init(PSRAM_PIO, PSRAM_SM, psram_send_pc, &g_psram_pio_cfg_tx);
    pio_sm_put_blocking(PSRAM_PIO, PSRAM_SM, (bytes * 8 / PSRAM_BITS) - 1);
    dma_channel_set_read_addr(g_psram_dma_send_chan, buffer, false);
    dma_channel_set_trans_count(g_psram_dma_send_chan, (bytes + sizeof(uint32_t) - 1) / sizeof(uint32_t), true);
    pio_sm_set_enabled(PSRAM_PIO, PSRAM_SM, true);
}

static void psram_send_nostop_start(void)
{
    pio_sm_set_enabled(PSRAM_PIO, PSRAM_SM, false);
    pio_sm_set_enabled(PSRAM_PIO, PSRAM_SM, false);
    dma_channel_abort(g_psram_dma_send_chan);
    dma_channel_abort(g_psram_dma_recv_chan);
    pio_sm_set_consecutive_pindirs(PSRAM_PIO, PSRAM_SM, PIN_SRAM_D0, PSRAM_BITS, true);
    pio_sm_init(PSRAM_PIO, PSRAM_SM, psram_send_nostop_pc, &g_psram_pio_cfg_tx);
    pio_sm_set_enabled(PSRAM_PIO, PSRAM_SM, true);
}

static void psram_send_nostop_stop(void)
{
    pio_sm_set_enabled(PSRAM_PIO, PSRAM_SM, false);
    pio_sm_set_enabled(PSRAM_PIO, PSRAM_SM, false);
    dma_channel_abort(g_psram_dma_send_chan);
    dma_channel_abort(g_psram_dma_recv_chan);
    pio_sm_set_consecutive_pindirs(PSRAM_PIO, PSRAM_SM, PIN_SRAM_D0, PSRAM_BITS, false);
}

static void psram_send_busy_wait(uint32_t *buffer, uint32_t bytes)
{
    psram_send(buffer, bytes);
    while (pio_interrupt_get(PSRAM_PIO, PSRAM_IRQ) == 0) {}
    pio_interrupt_clear(PSRAM_PIO, PSRAM_IRQ);
    pio_sm_set_enabled(PSRAM_PIO, PSRAM_SM, false);
    pio_sm_set_consecutive_pindirs(PSRAM_PIO, PSRAM_SM, PIN_SRAM_D0, PSRAM_BITS, false);
}

static void psram_recv(uint32_t *buffer, uint32_t dummy, uint32_t bytes)
{
    pio_sm_set_enabled(PSRAM_PIO, PSRAM_SM, false);
    dma_channel_abort(g_psram_dma_send_chan);
    dma_channel_abort(g_psram_dma_recv_chan);
    pio_sm_set_consecutive_pindirs(PSRAM_PIO, PSRAM_SM, PIN_SRAM_D0, PSRAM_BITS, false);
    pio_sm_init(PSRAM_PIO, PSRAM_SM, psram_recv_pc, &g_psram_pio_cfg_rx);
    pio_sm_put_blocking(PSRAM_PIO, PSRAM_SM, dummy);
    pio_sm_put_blocking(PSRAM_PIO, PSRAM_SM, (bytes * 8 / PSRAM_BITS) - 1);
    dma_channel_set_write_addr(g_psram_dma_recv_chan, buffer, false);
    dma_channel_set_trans_count(g_psram_dma_recv_chan, (bytes + sizeof(uint32_t) - 1) / sizeof(uint32_t), true);
    pio_sm_set_enabled(PSRAM_PIO, PSRAM_SM, true);
}

static void psram_recv_busy_wait(uint32_t *buffer, uint32_t dummy, uint32_t bytes)
{
    psram_recv(buffer, dummy, bytes);
    while (pio_interrupt_get(PSRAM_PIO, PSRAM_IRQ) == 0) {}
    pio_interrupt_clear(PSRAM_PIO, PSRAM_IRQ);
    pio_sm_set_enabled(PSRAM_PIO, PSRAM_SM, false);
}

static void psram_init_pin(void)
{
    pca9557_gpio_set(PCA_SRAM_CE, 1);

    gpio_init(PIN_SRAM_CLK);
    gpio_put(PIN_SRAM_CLK, 0);
    gpio_set_drive_strength(PIN_SRAM_CLK, GPIO_DRIVE_STRENGTH_12MA);
    gpio_set_slew_rate(PIN_SRAM_CLK, GPIO_SLEW_RATE_FAST);
    gpio_set_dir(PIN_SRAM_CLK, GPIO_OUT);

    busy_wait_us(150);

    pio_gpio_init(PSRAM_PIO, PIN_SRAM_CLK);

    for (int i = 0; i < PSRAM_BITS; i++) {
        gpio_init(PIN_SRAM_D0 + i);
        gpio_set_drive_strength(PIN_SRAM_D0 + i, GPIO_DRIVE_STRENGTH_12MA);
        pio_gpio_init(PSRAM_PIO, PIN_SRAM_D0 + i);
    }
    pio_sm_set_consecutive_pindirs(PSRAM_PIO, PSRAM_SM, PIN_SRAM_CLK, 1, true);
    pio_sm_set_consecutive_pindirs(PSRAM_PIO, PSRAM_SM, PIN_SRAM_D0, 8, false);
}

static void psram_init_pio_alloc(void)
{
    pio_sm_claim(PSRAM_PIO, PSRAM_SM);
    psram_send_pc = pio_add_program(PSRAM_PIO, &psram_send_program);
    psram_send_nostop_pc = pio_add_program(PSRAM_PIO, &psram_send_nostop_program);
    psram_recv_pc = pio_add_program(PSRAM_PIO, &psram_recv_program);
}

static void psram_init_pio_tx(void)
{
    g_psram_pio_cfg_tx = psram_send_program_get_default_config(psram_send_pc);
    sm_config_set_wrap(&g_psram_pio_cfg_tx, 0, 31);
    sm_config_set_clkdiv_int_frac(&g_psram_pio_cfg_rx, 1, 0);
    sm_config_set_in_shift(&g_psram_pio_cfg_tx, true, true, 32);
    sm_config_set_out_shift(&g_psram_pio_cfg_tx, true, true, 32);

    sm_config_set_in_pins(&g_psram_pio_cfg_tx, PIN_SRAM_D0);
    sm_config_set_out_pins(&g_psram_pio_cfg_tx, PIN_SRAM_D0, PSRAM_BITS);
    sm_config_set_set_pins(&g_psram_pio_cfg_tx, PIN_SRAM_D0, 0);
    sm_config_set_sideset_pins(&g_psram_pio_cfg_tx, PIN_SRAM_CLK);
    sm_config_set_fifo_join(&g_psram_pio_cfg_tx, PIO_FIFO_JOIN_TX);
}

static void psram_init_rx_hi_speed(bool high)
{ 
    uint32_t div_high = clock_get_hz(clk_sys) / (PSRAM_LIMIT_READ_HIGH_CLK_FREQ * 2);
    uint32_t div_low = clock_get_hz(clk_sys) / (PSRAM_LIMIT_READ_LOW_CLK_FREQ * 2);

    if (high) {
        sm_config_set_clkdiv_int_frac(&g_psram_pio_cfg_rx, div_high, 0);
    } else {
        sm_config_set_clkdiv_int_frac(&g_psram_pio_cfg_rx, div_low, 0);
    }
}

static void psram_init_pio_rx(void)
{
    uint32_t div = clock_get_hz(clk_sys) / (PSRAM_LIMIT_READ_HIGH_CLK_FREQ * 2);

    if (div == 0) {
        div = 1;
    }
    g_psram_pio_cfg_rx = psram_send_program_get_default_config(psram_recv_pc);
    sm_config_set_wrap(&g_psram_pio_cfg_rx, 0, 31);
    sm_config_set_clkdiv_int_frac(&g_psram_pio_cfg_rx, div, 0);
    sm_config_set_in_shift(&g_psram_pio_cfg_rx, true, true, 32);
    sm_config_set_out_shift(&g_psram_pio_cfg_rx, true, true, 32);

    sm_config_set_in_pins(&g_psram_pio_cfg_rx, PIN_SRAM_D0);
    sm_config_set_out_pins(&g_psram_pio_cfg_rx, PIN_SRAM_D0, PSRAM_BITS);
    sm_config_set_set_pins(&g_psram_pio_cfg_rx, PIN_SRAM_D0, 0);
    sm_config_set_sideset_pins(&g_psram_pio_cfg_rx, PIN_SRAM_CLK);
    sm_config_set_fifo_join(&g_psram_pio_cfg_rx, PIO_FIFO_JOIN_NONE);
}

static void psram_init_pio(void)
{
    psram_init_pio_alloc();
    psram_init_pio_tx();
    psram_init_pio_rx();
}

static void psram_init_dma(void)
{
    dma_channel_config cfg;

    g_psram_dma_send_chan = dma_claim_unused_channel(true);
    g_psram_dma_recv_chan = dma_claim_unused_channel(true);

    cfg = dma_channel_get_default_config(g_psram_dma_send_chan);
    channel_config_set_dreq(&cfg, PSRAM_DMA_TX_REQ);
    channel_config_set_transfer_data_size(&cfg, DMA_SIZE_32);
    channel_config_set_read_increment(&cfg, true);
    channel_config_set_write_increment(&cfg, false);
    channel_config_set_high_priority(&cfg, true);

    dma_channel_configure(
        g_psram_dma_send_chan,     // Channel to be configured
        &cfg,                // The configuration we just created
        &PSRAM_PIO->txf[PSRAM_SM], // The initial write address
        0,                       // The initial read address
        1,                       // Number of transfers; in this case each is 1 byte.
        false                    // Start immediately.
    );

    cfg = dma_channel_get_default_config(g_psram_dma_recv_chan);
    channel_config_set_dreq(&cfg, PSRAM_DMA_RX_REQ);
    channel_config_set_transfer_data_size(&cfg, DMA_SIZE_32);
    channel_config_set_read_increment(&cfg, false);
    channel_config_set_write_increment(&cfg, true);
    channel_config_set_high_priority(&cfg, true);

    dma_channel_configure(
        g_psram_dma_recv_chan,          // Channel to be configured
        &cfg,                       // The configuration we just created
        0,                              // The initial write address
        &PSRAM_PIO->rxf[PSRAM_SM],      // The initial read address
        1,                              // Number of transfers; in this case each is 1 byte.
        false                           // Start immediately.
    );
}

static void psram_init_cfg_qspi(void)
{
    uint32_t data35[] = {0x11110000, 0x11001100};

    pca9557_gpio_set(PCA_SRAM_CE, 0);
    psram_send_busy_wait(data35, sizeof(data35));
    pca9557_gpio_set(PCA_SRAM_CE, 1);
}

static void psram_init_cfg_spi(void)
{
    uint32_t dataf5[] = {0x55ff};

    pca9557_gpio_set(PCA_SRAM_CE, 0);
    psram_send_busy_wait(dataf5, 2);
    pca9557_gpio_set(PCA_SRAM_CE, 1);
}

static void psram_set_addr(uint32_t addr, uint32_t cmd)
{
    uint32_t write[2] = {0};
    uint32_t bit;

    addr = addr & PSRAM_ADDR_MASK;
    write[0] = cmd & 0xffff;
    bit = (addr >> 20) & 0xf;
    write[0] |= ((bit << 16) | (bit << 20));
    bit = (addr >> 16) & 0xf;
    write[0] |= ((bit << 24) | (bit << 28));

    bit = (addr >> 12) & 0xf;
    write[1] |= ((bit << 0) | (bit << 4));
    bit = (addr >> 8) & 0xf;
    write[1] |= ((bit << 8) | (bit << 12));
    bit = (addr >> 4) & 0xf;
    write[1] |= ((bit << 16) | (bit << 20));
    bit = (addr >> 0) & 0xf;
    write[1] |= ((bit << 24) | (bit << 28));

    psram_send_busy_wait(write, sizeof(write));
}

static void psram_write(uint32_t addr, uint32_t *value, uint32_t bytes)
{
    pca9557_gpio_set(PCA_SRAM_CE, 0);
    psram_set_addr(addr, PSRAM_WRITE_CMD);
    psram_send_busy_wait(value, bytes);
    pca9557_gpio_set(PCA_SRAM_CE, 1);
}

static void psram_nostop_write_start(uint32_t addr)
{
    pca9557_gpio_set(PCA_SRAM_CE, 0);
    psram_set_addr(addr, PSRAM_WRITE_CMD);
    psram_send_nostop_start();
}

static void psram_nostop_write_send(uint32_t data)
{
    pio_sm_put_blocking(PSRAM_PIO, PSRAM_SM, data);
}

static void psram_nostop_write_stop(void)
{
    psram_send_nostop_stop();
    pca9557_gpio_set(PCA_SRAM_CE, 1);
}

void __no_inline_not_in_flash_func(psram_read_start)(uint32_t addr, uint32_t *value, uint32_t bytes)
{
    uint32_t temp[2] = {0};

    bytes = (bytes + sizeof(uint32_t) - 1) & (~(sizeof(uint32_t) - 1));
    pca9557_gpio_set(PCA_SRAM_CE, 0);
    psram_set_addr(addr, PSRAM_READ_CMD);
    psram_recv(value, 6, bytes);
}

void __no_inline_not_in_flash_func(psram_read_wait)(void)
{
    while (pio_interrupt_get(PSRAM_PIO, PSRAM_IRQ) == 0) {}
    pio_interrupt_clear(PSRAM_PIO, PSRAM_IRQ);
    pio_sm_set_enabled(PSRAM_PIO, PSRAM_SM, false);
    pca9557_gpio_set(PCA_SRAM_CE, 1);
}

void __no_inline_not_in_flash_func(psram_read)(uint32_t addr, uint32_t *value, uint32_t bytes)
{
    uint32_t temp[2] = {0};

    bytes = (bytes + sizeof(uint32_t) - 1) & (~(sizeof(uint32_t) - 1));
    pca9557_gpio_set(PCA_SRAM_CE, 0);
    psram_set_addr(addr, PSRAM_READ_CMD);
    psram_recv_busy_wait(value, 6, bytes);
    pca9557_gpio_set(PCA_SRAM_CE, 1);
}

static uint64_t psram_read_die_id(void)
{
    uint32_t write[8] = {0};
    uint32_t read[18] = {0};
    uint32_t id = 0;

    write[0] = 0x11000011;
    write[1] = 0x11111111;
    write[2] = 0x0;
    write[3] = 0x0;
    write[4] = 0x0;
    write[5] = 0x0;
    write[6] = 0x0;
    write[7] = 0x0;
    pca9557_gpio_set(PCA_SRAM_CE, 0);
    psram_send_busy_wait(write, sizeof(write));
    psram_recv_busy_wait(read, 0, sizeof(read));
    pca9557_gpio_set(PCA_SRAM_CE, 1);
    busy_wait_us(100);

    for (int i = 0; i < 16; i++) {
        if (read[i / 4] & (0x2 << ((i % 4) * 8))) {
            id |= (1 << (31 - i));
        }
        if (read[i / 4] & (0x20 << ((i % 4) * 8))) {
            id |= (1 << (15 - i));
        }
    }
    return id;
}

static void psram_init_cfg_reset(void)
{
    uint32_t data66[] = {0x00111100, 0x00111100};
    uint32_t data99[] = {0x11000011, 0x11000011};

    pca9557_gpio_set(PCA_SRAM_CE, 0);
    psram_send_busy_wait(data66, sizeof(data66));
    pca9557_gpio_set(PCA_SRAM_CE, 1);
    busy_wait_us(200);
    pca9557_gpio_set(PCA_SRAM_CE, 0);
    psram_send_busy_wait(data99, sizeof(data99));
    pca9557_gpio_set(PCA_SRAM_CE, 1);
    busy_wait_us(1000);
}

static void psarm_init_test1(void)
{
    uint32_t write[8] = {0x12345678, 0x23456789, 0x34567890, 0x45678901, 0x56789012, 0x67890123, 0x78901234, 0x89012345};
    uint32_t read[8] = {0};

    psram_write(PSRAM_SIZE - 4, write, sizeof(write));
    // psram_write(0, write, sizeof(write));
    busy_wait_us(50);
    psram_read(0, read, sizeof(read));
    busy_wait_us(50);
}

static void psarm_init_test(void)
{
#if 1
    uint32_t len = psram_size() / 4;
    static uint32_t read[1024];


    psram_nostop_write_start(PSRAM_SIZE - 4);
    for (int i = 0; i < len; i++) {
        psram_nostop_write_send(0x5a5a5a5a);
    }
    busy_wait_us(50);
    psram_nostop_write_stop();
    busy_wait_us(50);
    for (int i = 0; i < len / 1024; i++) {
        psram_read(i * 1024 / 2, read, sizeof(read));
        for (int j = 0; j < 1024; j++) {
            if (read[j] != 0x5a5a5a5a) {
                log_printf("psram_test_check error\n");
                return;
            }
        }
    }    
    busy_wait_us(50);
    log_printf("psram_test_check done\n");
#endif
}

uint32_t psram_size(void)
{
    return PSRAM_SIZE;
}

void *psram_save_start(void)
{
    log_printf("psram_save_start\n");
    psram_nostop_write_start(0);
    return (void *)&PSRAM_PIO->txf[PSRAM_SM];
}

void psram_save_stop(void)
{
    // busy_wait_us(50);
    psram_nostop_write_stop();
    log_printf("psram_save_stop\n");
}

static void psarm_init_test_nostop(void)
{
    uint32_t read[4] = {0};
    uint32_t size = psram_size() / 4;
    
    psram_save_start();
    for (int i = 0; i < size; i++) {
        pio_sm_put_blocking(PSRAM_PIO, PSRAM_SM, 0x12345678);
    }
    for (int i = 0; i < size; i++) {
        pio_sm_put_blocking(PSRAM_PIO, PSRAM_SM, 0x87654321);
    }
    pio_sm_put_blocking(PSRAM_PIO, PSRAM_SM, 0x10203040);
    pio_sm_put_blocking(PSRAM_PIO, PSRAM_SM, 0x50607080);
    log_printf("write done\n");
    psram_save_stop();
    psram_read(0, read, sizeof(read));
    log_printf("psram_init_test_nostop read0 = 0x%x\n", read);
    psram_read(2, read, sizeof(read));
    log_printf("psram_init_test_nostop read1 = 0x%x\n", read);
    psram_read(4, read, sizeof(read));
    log_printf("psram_init_test_nostop read2 = 0x%x\n", read);
}

static void psram_init_cfg(void)
{
    psram_init_rx_hi_speed(false);

    psram_init_cfg_reset();
    do {
        psram_init_cfg_spi();
        psram_init_cfg_reset();
    } while (psram_read_die_id() != 0x0d5d0d5d);
    psram_init_cfg_reset();
    psram_read_die_id();
    psram_init_cfg_qspi();
    psram_init_rx_hi_speed(true);
    log_printf("psram_init_cfg done\n");
    // psarm_init_test();
    // psarm_init_test_nostop();
}

void psram_init(void)
{
    vreg_set_voltage(VREG_VOLTAGE_1_30);
    busy_wait_us(50 * 1000);
    set_sys_clock_hz(PSRAM_MAX_CPU_RATE, true);

    psram_init_pin();
    psram_init_pio();
    psram_init_dma();
    psram_init_cfg();
}

void psram_proc_set_cmd(char *value)
{
    uint32_t read = 0;

    log_printf("cmd: %s\n", value);
    if (!strcmp(value, "read")) {
        psram_read(0, &read, sizeof(read));
        log_printf("value: 0x%x\n", read);
    }
    if (!strcmp(value, "write0")) {
        psram_save_start();
        pio_sm_put_blocking(PSRAM_PIO, PSRAM_SM, 0x12346789);
        psram_save_stop();
    }
    if (!strcmp(value, "write1")) {
        psram_save_start();
        pio_sm_put_blocking(PSRAM_PIO, PSRAM_SM, 0x78787788);
        psram_save_stop();
    }
}
HTTPD_CGI_PARAM_DEFINE("Psram", psram_proc_set_cmd);
